package net.lucene.chapter11.highlighter;

import com.chenlb.mmseg4j.analysis.MMSegAnalyzer;
import org.apache.commons.io.FileUtils;
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.queryParser.MultiFieldQueryParser;
import org.apache.lucene.queryParser.QueryParser;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.highlight.*;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FSDirectory;
import org.apache.lucene.util.Version;
import org.apache.tika.Tika;

import java.io.File;
import java.io.IOException;

//@formatter:off
/**
 * 【第十一章】高亮
 * ------------------------------------------------------------
 *
 * 高亮功能属于Lucene的扩展功能（或者叫做贡献功能）
 * 其所需jar位于Lucene-3.6.2.zip中的/contrib/highlighter/文件夹中
 * 本例需要以下4个jar
 *
 * lucene-core-3.6.2.jar
 * lucene-highlighter-3.6.2.jar
 * mmseg4j-all-1.8.5-with-dic.jar
 * tika-app-1.4.jar
 *
 * @author Barudisshu
 */
//@formatter:on
public class HelloHighLighter {

    // 索引存储目录
    private static final String PATH_OF_INDEX = "chapter11-demo/src/main/resources/lucene/11_index";
    // 被索引文件目录
    private static final String PATH_OF_FILE = "chapter11-demo/src/main/resources/lucene/11_file";

    private Directory directory;
    private IndexReader reader;

    public HelloHighLighter() {

        Document doc;
        IndexWriter writer = null;

        try {

            directory = FSDirectory.open(new File(PATH_OF_INDEX));
            writer = new IndexWriter(directory, new IndexWriterConfig(Version.LUCENE_36, new MMSegAnalyzer()));
            writer.deleteAll();
            for (File file : FileUtils.listFiles(new File(PATH_OF_FILE), null, true)) {
                doc = new Document();
                doc.add(new Field("fileContent", new Tika().parse(file))); // Field.Store.NO,Field.Index.ANALYZED
                doc.add(new Field("filePath", file.getAbsolutePath(), Field.Store.YES, Field.Index.NOT_ANALYZED_NO_NORMS));
                writer.addDocument(doc);
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (null != writer) {
                try {
                    writer.close();
                } catch (IOException ce) {
                    ce.printStackTrace();
                }
            }
        }
    }

    /**
     * 获取IndexSearcher实例
     */
    private IndexSearcher getIndexSearcher() {
        try {
            if (reader == null) {
                reader = IndexReader.open(directory);
            } else {
                // if the index was changed since the provided reader was opened, open and return a new reader; else,return null
                // 如果当前reader在打开期间index发生改变,则打开并返回一个新的IndexReader,否则返回null
                IndexReader ir = IndexReader.openIfChanged(reader);
                if (ir != null) {
                    reader.close(); // 关闭原reader
                    reader = ir;    // 赋予新reader
                }
            }
            return new IndexSearcher(reader);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null; // 发生异常则返回null
    }


    /**
     * 高亮搜索
     *
     * @param expr 搜索表达式
     * @see 高亮搜索时,不建议把高亮信息存到索引里,而是搜索到内容之后再进行高亮处理
     * @see 这里用的是MMSeg4j中文分词器
     */
    public void searchByHighLighter(String expr) {
        Analyzer analyzer = new MMSegAnalyzer();
        IndexSearcher searcher = this.getIndexSearcher();
        // 搜索多个Field
        QueryParser parser = new MultiFieldQueryParser(Version.LUCENE_36, new String[]{"filePath", "fileContent"}, analyzer);
        try {
            Query query = parser.parse(expr);
            TopDocs tds = searcher.search(query, 50);
            for (ScoreDoc sd : tds.scoreDocs) {

                Document doc = searcher.doc(sd.doc);
                // 获取文档内容
                String fileContent = new Tika().parseToString(new File(doc.get("filePath"))).trim();
                System.out.println("搜索到的内容为：[" + fileContent + "]");
                // 开始高亮处理
                QueryScorer queryScorer = new QueryScorer(query);
                Fragmenter fragmenter = new SimpleSpanFragmenter(queryScorer, fileContent.length());
                Formatter formatter = new SimpleHTMLFormatter("<span style='color:red'>", "</span>");
                Highlighter hl = new Highlighter(formatter, queryScorer);
                hl.setTextFragmenter(fragmenter);
                System.out.println("高亮后的内容为：[" + hl.getBestFragment(analyzer, "fileContent", fileContent) + "]");
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (null != searcher) {
                try {
                    searcher.close(); // 记得关闭IndexSearcher
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }


    /**
     * 高亮的使用方式
     */
    private static void testHighLighter() {

        String fieldName = "myInfo"; //这个可以随便写,就是起个标识的作用
        String text = "我来自中国广东省高州市金山街道文龙村委下低屋村";
        QueryParser parser = new QueryParser(Version.LUCENE_36, fieldName, new MMSegAnalyzer());
        try {
            // MMSeg4j的new MMSegAnalyzer()默认只会对'中国'和'兴隆'进行分词,所以这里就只高亮它们俩了
            Query query = parser.parse("中国 金山");
            // 针对查询出来的文本,查询其评分,以便于能够根据评分决定显示情况
            QueryScorer queryScorer = new QueryScorer(query);
            // 对字符串或文本进行分段,SimpleSpanFragmenter构造方法的第二个参数可以指定高亮的文本长度,默认为100
            Fragmenter fragmenter = new SimpleSpanFragmenter(queryScorer);
            // 高亮时的高亮格式,默认为<B></B>,这里指定为红色字体
            Formatter formatter = new SimpleHTMLFormatter("<span style='color:red'>", "</span>");
            // Highlighter专门用来做高亮显示
            // 该构造方法还有一个参数为Encoder,它有两个实现类DefaultEncoder和SimpleHTMLEncoder
            // SimpleHTMLEncoder可以忽略掉HTML标签,而DefaultEncoder则不会忽略HTML标签
            Highlighter hl = new Highlighter(formatter, queryScorer);
            hl.setTextFragmenter(fragmenter);
            System.out.println(hl.getBestFragment(new MMSegAnalyzer(), fieldName, text));
        } catch (Exception e) {
            e.printStackTrace();
        }
    }


    /**
     * 小测试一下
     */
    public static void main(String[] args) {
        // 测试高亮的基本使用效果
        HelloHighLighter.testHighLighter();
        // 测试高亮搜索的效果(测试前记得在myExample/myFile/文件夹中准备一个或多个内容包含"依赖"的doc或pdf的等文件)
        new HelloHighLighter().searchByHighLighter("依赖");
    }
}
